package com.LeeCode;

import java.util.ArrayList;
import java.util.List;

/**
 * 删点成林
 */

public class Code1110 {
    public static void main(String[] args) {
        Integer[] arr = {1, 2, 3, 4, 5, 6, 7};
        TreeNode root = Utils.buildTree(arr);
        int[] nums = {3, 5};
        System.out.println(new Code1110().delNodes(root, nums));
    }

    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        List<TreeNode> ans = new ArrayList<>();
        int[] cnt = new int[1001];
        for (int num : to_delete)
            cnt[num]++;
        if (cnt[root.val] == 0) ans.add(root);
        dfs(root, ans, cnt);
        return ans;
    }

    private TreeNode dfs(TreeNode root, List<TreeNode> ans, int[] cnt) {
        if (root == null) return null;
        root.left = dfs(root.left, ans, cnt);
        root.right = dfs(root.right, ans, cnt);
        if (cnt[root.val] != 0) {
            if (root.left != null) ans.add(root.left);
            if (root.right != null) ans.add(root.right);
            return null;
        }
        return root;
    }
}
